"""

"""

import os
import uuid
from typing import Union
import pandas as pd
from datasets import load_dataset, Dataset, load_from_disk
from langsmith.schemas import Dataset
from prompt import industry_prompt_reason_template


def add_uuid(filename):
    assert os.path.exists(filename)
    df = pd.read_csv(filename)
    length = df.shape[0]

    if "uuid" not in df.columns:
        import uuid

        uuids = [str(uuid.uuid4()) for _ in range(length)]
        df["uuid"] = uuids
        df.to_csv(filename, index=False)


# 企业提示词构造
def get_industry_trans_func(attr_name, is_industry_prompt):
    """
    map func of dataset:
        dataset = dataset.map(build_prompt)
    """

    def _func(item):
        col_name = [
            "企业名称",
            "经营范围",
            "一级行业分类",
            "二级行业分类",
            "三级行业分类",
        ]
        industry_info = []
        pair_prompt = "* {key}: {value}"

        for col in col_name:
            industry_info.append(pair_prompt.format(key=col, value=item[col]))

        item[attr_name] = is_industry_prompt.format(
            industry_info="\n".join(industry_info)
        )
        return item

    return _func


def process_data(dataset: Dataset, prompt_reason_template) -> Dataset:
    dataset = dataset.map(
        get_industry_trans_func("industry_info", prompt_reason_template)
    )
    return dataset


def result_parse(dataset: Dataset, output_filename):
    """
    def get_last_folder(path):
    return os.path.basename(path.rstrip('/'))

    kws = ["半导体", "微电子", "电子元器件", "芯片", "设备和材料", "集成电路"]

    for kw in kws:
        folder = f"/home/jie/github/LLM/API/openai_examples/industry_dataset/data/csvs/宁波_{kw}"

        result_parse(
            folder,
            os.path.join(
                "/home/jie/github/LLM/API/openai_examples/industry_dataset/output",
                f"{get_last_folder(folder)}.csv"
            ),
        )
    """
    if isinstance(dataset, str):
        assert os.path.exists(dataset)
        dataset = load_from_disk(dataset)

    results = []

    for item in dataset:
        infer = item["infer"]
        rear = infer[-3:]

        if "是" in rear and "否" not in rear:
            cls_target = "是"
        elif "否" in rear and "是" not in rear:
            cls_target = "否"
        else:
            cls_target = rear
        results.append(cls_target)

    new_dataset = dataset.add_column("result", results)
    new_dataset = new_dataset.filter(lambda x: x["result"] == "是")
    new_dataset.to_csv(output_filename)


if __name__ == "__main__":
    # file = "/home/jie/github/LLM/API/openai_examples/industry_dataset/data/csvs/宁波_半导体.csv"
    # dataset = load_dataset("csv", data_files=file, split="train")
    # # print(df.head())
    # dataset = process_data(dataset)
    pass
